{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import re\n",
    "import sys\n",
    "import os\n",
    "from datetime import datetime\n",
    "import numpy as np\n",
    "from mpmath import mp\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axes_grid1 import host_subplot\n",
    "import mpl_toolkits.axisartist as AA\n",
    "import math\n",
    "import warnings\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "plt.rcParams['figure.dpi'] = 90\n",
    "plt.rcParams['figure.figsize'] = [24.0, 16.0]\n",
    "plt.rcParams['text.usetex'] = False\n",
    "plt.rcParams['mathtext.fontset'] = 'stixsans'\n",
    "plt.rcParams['font.family'] = 'DejaVu Sans'\n",
    "print(f'Current directory {os.getcwd()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CALL_GAS_CHUNK = 10000000\n",
    "NUM_TEST = 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def events(log):\n",
    "    EVENT_COST = 1954\n",
    "    start = None\n",
    "    discount = 0\n",
    "    events = []\n",
    "    for line in log.split('\\n'):\n",
    "        match = re.search(r'^(ENTER|LEAVE) (\\w+) (\\d+) (\\d+)$', line)\n",
    "        if match:\n",
    "            (event, name, gas, mem) = match.groups()\n",
    "            depth = (int(gas) - 1) // CALL_GAS_CHUNK\n",
    "            gas = -int(gas)\n",
    "            mem = int(mem)\n",
    "            if start:\n",
    "                gas -= start\n",
    "            else:\n",
    "                start = gas\n",
    "                gas = 0\n",
    "            begin = event == 'ENTER'\n",
    "            if not name in ['calldata', 'transaction']:\n",
    "                discount += EVENT_COST\n",
    "            gas -= discount\n",
    "            events += [(gas, begin, name, mem)]\n",
    "    return events"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(events):\n",
    "    stack = []\n",
    "    spans = []\n",
    "    alloc = []\n",
    "    histogram = {}\n",
    "    last = None\n",
    "    for event in events:\n",
    "        (time, begin, name, mem) = event\n",
    "        # Collect top of stack histogram\n",
    "        if not last is None:\n",
    "            top_of_stack = stack[-1][2]\n",
    "            duration = time - last\n",
    "            if top_of_stack in histogram:\n",
    "                histogram[top_of_stack] += duration\n",
    "            else:\n",
    "                histogram[top_of_stack] = duration\n",
    "        last = time\n",
    "        # Compute spans\n",
    "        if begin:\n",
    "            stack += [event]\n",
    "        else:\n",
    "            (start, _, previous_name, _) = stack[-1]\n",
    "            stack = stack[:-1]\n",
    "            assert name == previous_name\n",
    "            spans += [(len(stack), start, time, name)]\n",
    "        # Compute memory allocation\n",
    "        alloc += [[time, mem]]\n",
    "    histogram = {k: v for k, v in sorted(histogram.items(), key=lambda item: -item[1])}\n",
    "    # print(len(spans), len(histogram))\n",
    "    return spans, histogram, alloc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def flamegraph(spans):\n",
    "    labels = list(set([span[3] for span in spans]))\n",
    "    maxt = spans[-1][2]\n",
    "    colour_map = mpl.cm.rainbow(np.linspace(0, 1, len(labels)))\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    ax.set_xlim((0,maxt))\n",
    "    ax.set_ylim((0,7))\n",
    "\n",
    "    for span in spans:\n",
    "        (depth, start, end, label) = span\n",
    "        colour = colour_map[labels.index(label)]\n",
    "        rectangle = mpl.patches.Rectangle((start, depth), end - start, 1, fc=colour, ec='white')\n",
    "        ax.add_patch(rectangle)\n",
    "        if end - start > 0.05 * maxt:\n",
    "            ax.text(start + 0.005 * maxt, depth+0.05, label)\n",
    "        elif end - start > 0.01 * maxt:\n",
    "            ax.text(start + 0.005 * maxt, depth+0.05, label, rotation='vertical', verticalalignment='bottom')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = [process(events(open(f'gas-{i}.log').read())) for i in range(NUM_TEST)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = list(data[-1][1].keys())\n",
    "print(labels)\n",
    "values = np.array([[hist[label] for label in labels] for (_, hist, _) in data]).T\n",
    "fig, ax = plt.subplots()\n",
    "x = range(NUM_TEST)\n",
    "y = values # / np.sum(values, axis=0)\n",
    "ax.stackplot(x, y[::-1, :], labels=labels[::-1])\n",
    "ax.set_xlim((0,NUM_TEST))\n",
    "#ax.set_ylim((0,1))\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "ax.legend(handles[::-1], labels[::-1], loc='lower left')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INDEX = NUM_TEST - 1 # <-- change to view a different run\n",
    "\n",
    "flame, _, alloc = data[INDEX]\n",
    "alloc = np.array(alloc)\n",
    "flamegraph(flame)\n",
    "plt.plot(alloc[:, 0], 6.5 * alloc[:, 1] / max(alloc[:, 1]), color='black')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
